#!/bin/bash

# 简单版URPC2020训练脚本

echo "======================================================================"
echo "🌊 URPC2020水下目标检测训练"
echo "======================================================================"
echo ""

# 检查数据集
echo "📊 检查URPC2020数据集..."
if [ ! -d "datasets/urpc2020" ]; then
    echo "❌ 数据集目录不存在: datasets/urpc2020"
    exit 1
fi

if [ ! -f "urpc2020.yaml" ]; then
    echo "❌ 数据集配置文件不存在: urpc2020.yaml"
    exit 1
fi

# 统计数据集
TRAIN_COUNT=$(ls -1 datasets/urpc2020/train/images/*.jpg 2>/dev/null | wc -l)
VALID_COUNT=$(ls -1 datasets/urpc2020/valid/images/*.jpg 2>/dev/null | wc -l)
TEST_COUNT=$(ls -1 datasets/urpc2020/test/images/*.jpg 2>/dev/null | wc -l)

echo "✅ 数据集已就绪:"
echo "   训练集: $TRAIN_COUNT 张图像"
echo "   验证集: $VALID_COUNT 张图像"
echo "   测试集: $TEST_COUNT 张图像"
echo ""

# 检查GPU
echo "🖥️  检查GPU状态..."
nvidia-smi --query-gpu=index,name,memory.total,memory.free --format=csv,noheader,nounits | while read line; do
    echo "   GPU: $line"
done
echo ""

# 训练参数
MODEL="yolov8n.pt"
DATA="urpc2020.yaml"
EPOCHS=120
BATCH=16
IMGSZ=640
DEVICE="0"

echo "⚙️  训练参数:"
echo "   模型: $MODEL"
echo "   数据: $DATA"
echo "   轮数: $EPOCHS epochs"
echo "   批次: $BATCH"
echo "   尺寸: ${IMGSZ}px"
echo "   设备: GPU $DEVICE"
echo ""

# 预估训练时间
echo "⏱️  预估训练时间:"
echo "   基于RTX 4090, 约1.5-2.0小时"
echo ""

# 确认开始
echo "======================================================================"
echo "🚀 准备开始训练..."
echo "======================================================================"
echo ""

read -p "按Enter键开始训练，或Ctrl+C取消..." -t 10 || echo ""

echo ""
echo "======================================================================"
echo "🏋️ 开始训练..."
echo "======================================================================"
echo ""

# 记录开始时间
START_TIME=$(date +%s)

# 运行训练
python train_urpc2020_simple.py \
    --model $MODEL \
    --data $DATA \
    --epochs $EPOCHS \
    --batch $BATCH \
    --imgsz $IMGSZ \
    --device $DEVICE \
    --project runs/train \
    --name urpc2020_yolov8 \
    --exist-ok

# 记录结束时间
END_TIME=$(date +%s)
DURATION=$((END_TIME - START_TIME))
HOURS=$((DURATION / 3600))
MINUTES=$(((DURATION % 3600) / 60))

echo ""
echo "======================================================================"
echo "✅ 训练脚本执行完成"
echo "======================================================================"
echo ""
echo "⏱️  训练用时: ${HOURS}小时 ${MINUTES}分钟"
echo ""
echo "📝 查看结果:"
echo "   训练目录: runs/train/urpc2020_yolov8/"
echo "   最佳模型: runs/train/urpc2020_yolov8/weights/best.pt"
echo "   最终模型: runs/train/urpc2020_yolov8/weights/last.pt"
echo ""
echo "🔬 进一步验证:"
echo "   1. 测试集验证: yolo val model=runs/train/urpc2020_yolov8/weights/best.pt data=urpc2020.yaml"
echo "   2. 可视化预测: yolo predict model=runs/train/urpc2020_yolov8/weights/best.pt source=datasets/urpc2020/test/images"
echo ""
